#!/usr/bin/python

__description__ = 'rawdata_replica'
__created__ = '2019/01/28'
__author__ = 'yangxiao'
__version__ = 'v1.0'

from tqdm import tqdm
from pymongo import MongoClient
from argparse import ArgumentParser
import sys
import time

parser = ArgumentParser(description=__description__)
parser.add_argument('-v', '--version', action='version', version=' '.join((__version__, __author__, __created__)))
parser.add_argument('--host', default='127.0.0.1', help='mongoDB server host, default is 127.0.0.1')
parser.add_argument('-p', '--port', default=27017, type=int, help='mongoDB server port, default is 27017')
parser.add_argument('database', help='mongoDB database name')
parser.add_argument('collection', help='mongoDB collection name')

parser.add_argument('id', help='download batch number, such as dlst_0000001')
parser.add_argument('downloader', help='downloader name')
parser.add_argument('-s', '--listSize', default=2000, type=int, help='out size, default is 2000')

BAT_TEMPLATE = """

@echo off
set "msg=download failed, retry"

rem running aspera downloader
echo start downloading
".\Aspera\Aspera Connect\bin\ascp" -i ".\Aspera\Aspera Connect\etc\asperaweb_id_dsa.openssh" -k1 -QTr -l180m --user=anonftp --host=ftp.ncbi.nlm.nih.gov --mode=recv --file-pair-list={file_pair_list} --file-manifest=text --file-manifest-path=log_mani -L log .

if "%errorlevel%"=="0" ( set "msg=download finished")

echo %msg%

mshta vbscript:msgbox(%0"%msg%",64,%0"download")(window.close)

"""


def main(host: str, port: int, database: str, collection: str,
         id_: str, downloader: str, list_size: int):
    if list_size <= 0: raise ValueError(list_size)

    today = lambda: time.strftime('%Y-%m-%d', time.localtime())
    download_url = lambda x: 'anonftp@ftp.ncbi.nlm.nih.gov:' \
                             '/sra/sra-instant/reads/ByRun/sra/{}/{}/{}/{}.sra'.format(x[:3], x[:6], x, x)

    client = MongoClient(host, port)
    collection = client[database][collection]
    query_filter = {
        'data_store_path': {'$exists': False},
        'downloader': {'$exists': False},
        '$or': [
            {'download_count': {'$exists': False}},
            {'download_count': {'$lte': 3}}
        ],
        'source': 'SRA'
    }
    projection = {
        'RunID': 1,
        'download_count': 1
    }
    cursor = collection.find(filter=query_filter, projection=projection).limit(list_size)
    total = cursor.count(with_limit_and_skip=True)

    if total == 0:
        client.close()
        print('No data output.')
        sys.exit(0)

    path_file = open(id_, 'w', encoding='utf-8')
    print('Output:', path_file.name)

    for it in tqdm(cursor, total=total):
        run_id = it['RunID']
        # write download way
        path_file.write('/sra/sra-instant/reads/ByRun/sra/{}/{}/{}/{}.sra\n'.format(
            run_id[:3], run_id[:6], run_id, run_id))
        # write save way
        path_file.write('{}/{}.sra\n'.format(id_, run_id))

        update = {
            '$set': {
                'downloader': downloader,
                'Yao_to_downloader': today(),
                'download_count': 1 if 'download_count' not in it else it['download_count'] + 1,
                'data_download_url': download_url(run_id)
            },
            '$addToSet': {
                'list_ids': id_
            }
        }
        collection.update_one({'_id': it['_id']}, update=update)

    path_file.close()

    bat_file = open(id_ + '.bat', 'w', encoding='utf-8')
    print('Output:', bat_file.name)
    bat_file.write(BAT_TEMPLATE.format(file_pair_list=id_))
    bat_file.close()

    client.close()


if __name__ == '__main__':
    args = parser.parse_args()
    main(args.host, args.port, args.database, args.collection, args.id, args.downloader, args.listSize)
    print('Done!')
